from graphviz import Digraph
from datetime import datetime

def create_detailed_model_diagram():
    diagram = Digraph("EfficientNetV2_GCA", format="png", engine="dot")
    diagram.attr(rankdir="TB", fontsize="10")  # 修改这里，将LR改为TB

    # Input
    diagram.node("Input", "Input Image\n224x224x3", shape="rect", style="filled", color="lightblue")

    # Stem Block
    diagram.node("Stem", "Stem Block\nConv2D (3x3, stride=2, 32 filters)\nBN + SiLU\n112x112x32",
                 shape="rect", style="filled", color="lightyellow")
    diagram.edge("Input", "Stem")

    # EfficientNetV2 Block + GCA Block Loop
    prev = "Stem"
    blocks = [
        ("Eff_Block_1", "EfficientNetV2 Block 1\nDepthwise Conv + SE\nStride=1, 112x112x32", "lightgreen"),
        ("GCA_Block_1", "GCA Block 1\nGrouped Channel Attention\n112x112x32", "pink"),
        ("Eff_Block_2", "EfficientNetV2 Block 2\nDepthwise Conv + SE\nStride=2, 56x56x64", "lightgreen"),
        ("GCA_Block_2", "GCA Block 2\nGrouped Channel Attention\n56x56x64", "pink"),
        ("Eff_Block_3", "EfficientNetV2 Block 3\nDepthwise Conv + SE\nStride=2, 28x28x128", "lightgreen"),
        ("GCA_Block_3", "GCA Block 3\nGrouped Channel Attention\n28x28x128", "pink"),
        ("Eff_Block_4", "EfficientNetV2 Block 4\nDepthwise Conv + SE\nStride=2, 14x14x256", "lightgreen"),
        ("GCA_Block_4", "GCA Block 4\nGrouped Channel Attention\n14x14x256", "pink"),
    ]
    for name, label, color in blocks:
        diagram.node(name, label, shape="rect", style="filled", color=color)
        diagram.edge(prev, name)
        prev = name

    # GAP
    diagram.node("GAP", "Global Avg Pooling\n1x1x256", shape="rect", style="filled", color="lightgray")
    diagram.edge(prev, "GAP")

    # Dropout
    diagram.node("Dropout", "Dropout (rate=0.5)\n1x1x256", shape="rect", style="filled", color="orange")
    diagram.edge("GAP", "Dropout")

    # Fully Connected
    diagram.node("FC", "Fully Connected\nSoftmax, 5 classes\n1x1x5", shape="rect", style="filled", color="purple")
    diagram.edge("Dropout", "FC")

    # Render
    now = datetime.now()
    now_str = now.strftime("%Y%m%d_%H%M%S")
    diagram.render("Detailed_EfficientNetV2_GCA"+now_str, cleanup=True)
    print("Model diagram saved as 'Detailed_EfficientNetV2_GCA.png'")


create_detailed_model_diagram()
